# @File : hello.py
# @Author : xudu
# @Time : 2024/11/19

import codon

# 读取FASTA文件
def read_fasta_file(file_path):
    result_dict = {}
    current_key = None
    current_value = []

    with open(file_path, 'r', encoding='utf-8') as file:
        for line in file:
            line = line.strip()
            # 跳过空行
            if not line:
                continue
            if line.startswith('>'):
                if current_key is not None:
                    result_dict[current_key] = ''.join(current_value)
                current_key = line[1:]
                current_value = []
            else:
                current_value.append(line)

        # 处理最后一个结果
        if current_key is not None:
            result_dict[current_key] = ''.join(current_value)

    return result_dict

# ATCG互换函数
def complement_dna(dna):
    complement_dict = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C'}
    result = ''.join(complement_dict.get(base, base) for base in dna)
    return result

# 查找开放阅读框
# 1. 将每个DNA序列通缩012索引分为三份
# 2. 将每个序列进行互补
# 3. 将得到的24份DNA序列存储在list中


# dna_sequences为DNA序列的list
def find_open_reading_frames(dna_sequences):
    result_list = []
    temp1 = []
    for dna in dna_sequences:
        for j in range(3):
            temp1.append(dna[j:])

    for i in range(len(temp1)):
        seq = temp1[i]
        result = complement_dna(seq)
        result = result[::-1]
        result_list.append(result)

    result_list += temp1
    print("打印可能的阅读框")
    for i in range(len(result_list)):
        print(i , ':',result_list[i])



    map = {}

    a = -1
    b = -1
    tempList = []
    for i in range(len(result_list)):
        seq = result_list[i]
        for j in range(0, len(seq)):
            # temp_seq = ''
            if seq[j:j + 3] == 'ATG':
                tempseq = seq[j:]
                tempList.append(tempseq)
                break
                # print(j, '=========:', tempseq)
                # for k in range(j, len(tempseq), 3):
                #
                #     if tempseq[k:k + 3] == 'TAA' or tempseq[k:k + 3] == 'TAG' or tempseq[k:k + 3] == 'TGA':
                #         b = k
                #         temp_seq = seq[:b]
                #         map[i] = temp_seq
    for i in range(len(tempList)):
        tseq = tempList[i]
        for j in range(0, len(tseq), 3):
            if tseq[j:j + 3] == 'TAA' or tseq[j:j + 3] == 'TAG' or tseq[j:j + 3] == 'TGA':
                temp_seq = tseq[:j]
                map[i] = temp_seq
                break

    return map

# 将DNA序列翻译为蛋白质 codon_dict为翻译表

def translate_dna(dna, codon_dict):
    result_list = []
    for i in range(0, len(dna), 3):
        if dna[i:i + 3] in codon_dict:
            result_list.append(codon_dict[dna[i:i + 3]])
    return ''.join(result_list)


if __name__ == '__main__':
    # 读取文件
    file_path = 'lab_results.txt'
    result_dict = read_fasta_file(file_path)
    # for key, value in result_dict.items():
    #     print(key, value)
    # 输出原始文件
    print('原始文件')
    values_list = list(result_dict.values())
    for i in range(len(values_list)):
        print(values_list[i])

    #=======================================================================
    # 查找开放阅读框
    # 1. 将每个DNA序列通缩012索引分为三份
    # 2. 将每个序列进行互补
    # 3. 将得到的24份DNA序列存储在list中
    # 4. 通过翻译表翻译出所有可能的orf

    map = find_open_reading_frames(values_list)

    # 输出可能的开放阅读框
    print('开放阅读框')
    for key, value in map.items():
        print(key, value)

    # 导入翻译表
    codon_dict = codon.codonDict

    # 翻译
    print('翻译')
    for key, value in map.items():
        print(key, translate_dna(value, codon_dict))

    # 读取感兴趣的蛋白质
    path = 'proteins_of_interest.txt'
    proteins = read_fasta_file(path)
    print('毒理学家提取的蛋白质数据')
    proteins_list = list(proteins.values())
    for i in range(len(proteins_list)):
        print(proteins_list[i])

    print("眼镜王蛇的蛋白质")
    list_ = proteins_list[3]
    print(list_)

    for key, value in map.items():
        danbaizhi = translate_dna(value, codon_dict)
        if danbaizhi == list_:
            print("嫌疑人有罪！")


    print("嫌疑人无罪！！！！！！！！！！！！！！！！")